import os
import mindspore.dataset as ds
import mindspore.dataset.engine as de #
import mindspore.dataset.vision.c_transforms as c
import mindspore.dataset.transforms.c_transforms as C2
from mindspore.communication.management import init, get_rank, get_group_size
import mindspore.common.dtype as mstype

def create_dataset(dataroot, batchSize=32, imageSize=32, repeat_num=1, workers=8, target='Ascend'):
    """Create dataset"""
    rank_id = 0
    device_num = 1

    # define map operations
    resize_op = c.Resize(imageSize)
    center_crop_op = c.CenterCrop(imageSize)
    normalize_op = c.Normalize(mean=(0.5*255, 0.5*255, 0.5*255), std=(0.5*255, 0.5*255, 0.5*255))
    hwc2chw_op = c.HWC2CHW()
    data_set = ds.Cifar10Dataset(dataroot, num_parallel_workers=workers, shuffle=True)
    transform = [resize_op, normalize_op, hwc2chw_op]

    type_cast_op = C2.TypeCast(mstype.int32)

    data_set = data_set.map(input_columns='image', operations=transform, num_parallel_workers=workers)
    data_set = data_set.map(input_columns='label', operations=type_cast_op, num_parallel_workers=workers)

    data_set = data_set.batch(batchSize, drop_remainder=True)
    data_set = data_set.repeat(repeat_num)

    return data_set


def create_dataset1(dataset_path, do_train, repeat_num=1, batch_size=32, target="Ascend", distribute=False,
                    enable_cache=False, cache_session_id=None):
    """
    create a train or evaluate cifar10 dataset for resnet50
    Args:
        dataset_path(string): the path of dataset.
        do_train(bool): whether dataset is used for train or eval.
        repeat_num(int): the repeat times of dataset. Default: 1
        batch_size(int): the batch size of dataset. Default: 32
        target(str): the device target. Default: Ascend
        distribute(bool): data for distribute or not. Default: False
        enable_cache(bool): whether tensor caching service is used for eval. Default: False
        cache_session_id(int): If enable_cache, cache session_id need to be provided. Default: None

    Returns:
        dataset
    """
    if distribute:
        init()
        #rank_id = get_rank()
        #device_num = get_group_size()
        rank_id = int(os.getenv('DEVICE_ID')) #
        device_num = int(os.getenv('RANK_SIZE')) #
    else:
        device_num = 1
    ds.config.set_prefetch_size(64)

    if device_num == 1:
        #data_set = ds.Cifar10Dataset(dataset_path, num_parallel_workers=12, shuffle=True)
        data_set = de.Cifar10Dataset(dataset_path, num_parallel_workers=8, shuffle=True) #
    else:
        #data_set = ds.Cifar10Dataset(dataset_path, num_parallel_workers=12, shuffle=True,
                           #num_shards=device_num, shard_id=rank_id)
        data_set = de.Cifar10Dataset(dataset_path, num_parallel_workers=8, shuffle=True,
                           num_shards=device_num, shard_id=rank_id) #
    # define map operations
    trans = []
    if do_train:
        trans += [
            c.RandomCrop((32, 32), (4, 4, 4, 4)),
            c.RandomHorizontalFlip(prob=0.5)
        ]

    trans += [
        c.Rescale(1.0 / 255.0, 0.0),
        c.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
        c.HWC2CHW()
    ]

    type_cast_op = C2.TypeCast(mstype.int32)

    data_set = data_set.map(operations=type_cast_op, input_columns="label", num_parallel_workers=8)
    # only enable cache for eval
    if do_train:
        enable_cache = False
    if enable_cache:
        if not cache_session_id:
            raise ValueError("A cache session_id must be provided to use cache.")
        eval_cache = ds.DatasetCache(session_id=int(cache_session_id), size=0)
        data_set = data_set.map(operations=trans, input_columns="image", num_parallel_workers=8, cache=eval_cache)
    else:
        data_set = data_set.map(operations=trans, input_columns="image", num_parallel_workers=8)

    # apply batch operations
    data_set = data_set.batch(batch_size, drop_remainder=True)
    # apply dataset repeat operation
    data_set = data_set.repeat(repeat_num)

    return data_set